from typing import Optional
from dataclasses import dataclass, field

@dataclass
class ModelConfig:
    model_name_or_path: str = field(
        default="MODEL_NAME",
        metadata={"help": "预训练模型路径或 Hugging Face Model Hub 模型标识符"}
    )
    model_revision: Optional[str] = field(
        default="main",
        metadata={"help": "指定模型版本 (分支名, tag 名 或 commit id)"}
    )
    torch_dtype: Optional[str] = field(
        default="bfloat16",
        metadata={"help": "覆盖默认 torch_dtype 以指定 dtype 加载模型"}
    )
    trust_remote_code: bool = field(
        default=True,
        metadata={"help": "加载模型和 tokenizer 时, 信任远程代码"}
    )
    attn_implementation: Optional[str] = field(
        default="flash_attention_2",
        metadata={"help": "选择 Attention 实现方式, 可选 'flash_attention_2' 或 None"}
    )
